/**
 * @author Gao Jian
 * @email gaoj@njupt.edu.cn
 * @create date 2021-06-03 10:23:40
 * @desc
 */

#include "cuda.h"
#include "defs.h"

__global__ void histogram_i(unsigned *hist_i, unsigned *hist_o,
                            unsigned char *data, unsigned char *mark,
                            unsigned n) {
  unsigned nz = n;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;
    if (mark[idx])
      atomicAdd(hist_i + data[idx], 1);
    else
      atomicAdd(hist_o + data[idx], 1);
  }
}

void make_histogram(unsigned *hist_i, unsigned *hist_o, unsigned char *data,
                    unsigned char *mark, unsigned nh, unsigned nd) {
  unsigned *d_hist_i = NULL;
  unsigned *d_hist_o = NULL;
  unsigned char *d_data = NULL;
  unsigned char *d_mark = NULL;
  cudaMalloc((void **)&d_hist_i, nh * sizeof(unsigned));
  cudaMalloc((void **)&d_hist_o, nh * sizeof(unsigned));
  cudaMalloc((void **)&d_data, nd * sizeof(unsigned char));
  cudaMalloc((void **)&d_mark, nd * sizeof(unsigned char));

  cudaMemcpy(d_hist_i, hist_i, nh * sizeof(unsigned), cudaMemcpyHostToDevice);
  cudaMemcpy(d_hist_o, hist_o, nh * sizeof(unsigned), cudaMemcpyHostToDevice);
  cudaMemcpy(d_data, data, nd * sizeof(unsigned char), cudaMemcpyHostToDevice);
  cudaMemcpy(d_mark, mark, nd * sizeof(unsigned char), cudaMemcpyHostToDevice);

  histogram_i<<<NUM_BLOCKS, NUM_THREADS>>>(d_hist_i, d_hist_o, d_data, d_mark,
                                           nd);

  cudaMemcpy(hist_i, d_hist_i, nh * sizeof(unsigned), cudaMemcpyDeviceToHost);
  cudaMemcpy(hist_o, d_hist_o, nh * sizeof(unsigned), cudaMemcpyDeviceToHost);

  for (int t = 1; t >= 0; --t) {
    hist_i[t] += hist_i[t + 1];
    hist_o[t] += hist_o[t + 1];
  }

  cudaFree(d_mark);
  cudaFree(d_data);
  cudaFree(d_hist_o);
  cudaFree(d_hist_i);
}

__global__ void vind_i(float *vind, unsigned char *datav, unsigned char *dataw,
                       unsigned char *datab, unsigned char *datar, float ratio,
                       unsigned sv, unsigned sw, unsigned sb, unsigned sz) {
  unsigned nz = (sv - 1) * (sw - 1) * (sb - 1);
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;
    unsigned ib = idx;
    unsigned iv = ib % (sv - 1);
    ib /= (sv - 1);
    unsigned iw = ib % (sw - 1);
    ib /= (sw - 1);

    unsigned f1 = 0, f2 = 0;
    for (unsigned t = 0; t < sz; t++) {
      if (datav[t] > iv && dataw[t] > iw && datab[t] > ib) {
        if (datar[t] == 0)
          f2++;
      } else {
        if (datar[t] == 1)
          f1++;
      }
    }
    vind[idx] = (fabs(ratio * f1 - f2) + (ratio * f1 + f2)) / sz;
    // if (idx == 0)
    //   printf("  vind[%d] = %f  \n", idx, vind[idx]);
  }
}

// __global__ void vind_i2(float *vind, unsigned char *datav, unsigned char
// *dataw,
//                        unsigned char *datab, unsigned char *datar, float
//                        ratio, unsigned *f1, unsigned *f2, unsigned icx,
//                        unsigned sv, unsigned sw, unsigned sb, unsigned sz) {
//   unsigned nz = sz;
//   unsigned increment = blockDim.x * gridDim.x;
//   unsigned nits = nz / increment + 1;
//   unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

//   unsigned ib = icx;
//   unsigned iv = ib % (sv - 1);
//   ib /= (sv - 1);
//   unsigned iw = ib % (sw - 1);
//   ib /= (sw - 1);

//   for (unsigned i = 0; i < nits; i++, idx += increment) {
//     if (idx >= nz)
//       break;

//     if (datav[idx] > iv && dataw[idx] > iw && datab[idx] > ib) {
//       if (datar[idx] == 0)
//         atomicAdd(f2, 1);
//     } else {
//       if (datar[idx] == 1)
//         atomicAdd(f1, 1);
//     }

//     vind[icx] = (fabs(ratio * *f1 - *f2) + (ratio * *f1 + *f2)) / sz;
//   }
// }

__global__ void vind_i3(unsigned *area, unsigned *vind_i, unsigned *vind_o,
                        unsigned char *datav, unsigned char *dataw,
                        unsigned char *datab, unsigned char *datar, unsigned sv,
                        unsigned sw, unsigned sb, unsigned sz) {
  unsigned nz = sz;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    unsigned icx = (datab[idx] * sw + dataw[idx]) * sv + datav[idx];

    if (datar[idx] == 1) {
      atomicAdd(vind_i + icx, 1);
      atomicAdd(area, 1);
    } else
      atomicAdd(vind_o + icx, 1);
    // printf("founf (%d) %d - [%d/%d, %d/%d, %d/%d] .\n", icx, datar[idx],
    //        datav[idx], sv, dataw[idx], sw, datab[idx], sb);
  }
}

__global__ void flip_vind(unsigned *vind_i, unsigned *vind_o, unsigned ni,
                          unsigned di, unsigned dbx, unsigned dby) {
  unsigned ix = threadIdx.x;
  unsigned bx = blockIdx.x;
  unsigned by = blockIdx.y;
  __shared__ unsigned vi[255], vo[255];

  // unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;
  // printf("pvi[%d] = %d \n", idx, vind_i[idx]);

  unsigned *pvi = vind_i + bx * dbx + by * dby;
  unsigned *pvo = vind_o + bx * dbx + by * dby;
  vi[ix] = pvi[ix * di];
  vo[ix] = pvo[ix * di];
  __syncthreads();
  // unsigned vi = 0;
  // unsigned vo = 0;

  unsigned id;
  for (id = 2; id <= 256; id *= 2) {
    if (ix % id == 0 && ix + id / 2 < ni) {
      vi[ix] += vi[ix + id / 2];
      vo[ix] += vo[ix + id / 2];
    }
    __syncthreads();
  }

  unsigned vli = vi[ix];
  unsigned vlo = vo[ix];
  unsigned ib = 0;
  for (id = 256; id > 1; id /= 2) {
    if (ix % id == 0)
      break;
    unsigned dt = id / 2;
    if ((ix % id) / dt == 0) {
      if (ib + dt < ni) {
        vli += vi[ib + id];
        vlo += vo[ib + id];
      }
    } else {
      ib += dt;
    }
    __syncthreads();
  }

  pvi[ix * di] = vli;
  pvo[ix * di] = vlo;

  printf("pvi[%d, %d, %d] = %d \n", ix, bx, by, vli);
}

__global__ void make_vindf(float *vind_f, unsigned *vind_i, unsigned *vind_o,
                           float ratio, unsigned sr, unsigned sz) {
  unsigned nz = sz;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;
    unsigned f1 = sr - vind_i[idx];
    vind_f[idx] =
        (fabs(ratio * f1 - vind_o[idx]) + (ratio * f1 + vind_o[idx])) / sr;

    // if (idx < 100)
    printf("vind_f[%d] = %f \n", idx, vind_f[idx]);
  }
}

__global__ void min_ii(unsigned *idx, float *vind, unsigned sz) {

  float vmin = vind[0];
  unsigned imin = 0;
  printf("  PPP  \n");
  for (unsigned i = 1; i < sz; i++) {
    if (vind[i] < vmin) {
      vmin = vind[i];
      imin = i;
    }
    if (i == 0)
      printf("  PPP[%d] = %f  \n", i, vind[i]);
  }
  *idx = imin;
  printf("  SSS  \n");
}

void make_vind(unsigned char *datao, unsigned char *datav, unsigned char *dataw,
               unsigned char *datab, unsigned char *datar, float ratio,
               unsigned sv, unsigned sw, unsigned sb, unsigned sz,
               unsigned char verbose) {
  unsigned *d_vind_i = NULL;
  unsigned *d_vind_o = NULL;
  float *d_vind_f = NULL;
  // unsigned char *d_buff = NULL;
  unsigned char *d_datav = NULL;
  unsigned char *d_dataw = NULL;
  unsigned char *d_datab = NULL;
  unsigned char *d_datar = NULL;
  unsigned ss = sv * sw * sb; //(sv - 1) * (sw - 1) * (sb - 1);
  unsigned szc = sz * sizeof(unsigned char);
  cudaMalloc((void **)&d_vind_i, ss * sizeof(unsigned));
  cudaMalloc((void **)&d_vind_o, ss * sizeof(unsigned));
  cudaMalloc((void **)&d_vind_f, ss * sizeof(float));
  // cudaMalloc((void **)&d_buff, ss * sizeof(float));
  cudaMalloc((void **)&d_datav, szc);
  cudaMalloc((void **)&d_dataw, szc);
  cudaMalloc((void **)&d_datab, szc);
  cudaMalloc((void **)&d_datar, szc);
  float *vind = new float[ss];

  cudaMemcpy(d_datav, datav, szc, cudaMemcpyHostToDevice);
  cudaMemcpy(d_dataw, dataw, szc, cudaMemcpyHostToDevice);
  cudaMemcpy(d_datab, datab, szc, cudaMemcpyHostToDevice);
  cudaMemcpy(d_datar, datar, szc, cudaMemcpyHostToDevice);

  // for (unsigned i = 0; i < ss; i++) {
  //   // cudaMemcpy(f1, &idmin, sizeof(unsigned), cudaMemcpyHostToDevice);
  //   // cudaMemcpy(f2, &idmin, sizeof(unsigned), cudaMemcpyHostToDevice);
  //   // vind_i<<<NUM_BLOCKS, NUM_THREADS>>>(d_vind, d_datav, d_dataw, d_datab,
  //   //                                     d_datar, ratio, f1, f2, i, sv, sw,
  //   //                                     sb, sz);
  //   if (i % ((sv - 1)) == 0)
  //     printf("+");
  // }

  // vind_i<<<NUM_BLOCKS, NUM_THREADS>>>(d_vind, d_datav, d_dataw, d_datab,
  //                                     d_datar, ratio, sv, sw, sb, sz);
  nppsSet_32u(0, d_vind_i, ss);
  nppsSet_32u(0, d_vind_o, ss);
  unsigned area = 0;
  unsigned *parea = (unsigned *)d_vind_f;
  cudaMemcpy(parea, &area, sizeof(unsigned), cudaMemcpyHostToDevice);
  vind_i3<<<NUM_BLOCKS, NUM_THREADS>>>(parea, d_vind_i, d_vind_o, d_datav,
                                       d_dataw, d_datab, d_datar, sv, sw, sb,
                                       sz);
  cudaMemcpy(&area, parea, sizeof(unsigned), cudaMemcpyDeviceToHost);

  printf("AREA = %d \n", area);

  flip_vind<<<dim3(sw, sb), sv>>>(d_vind_i, d_vind_o, sv, 1, sv, sv * sw);
  flip_vind<<<dim3(sv, sb), sw>>>(d_vind_i, d_vind_o, sw, sv, 1, sv * sw);
  flip_vind<<<dim3(sv, sw), sb>>>(d_vind_i, d_vind_o, sb, sv * sw, 1, sv);

  make_vindf<<<NUM_BLOCKS, NUM_THREADS>>>(d_vind_f, d_vind_i, d_vind_o, ratio,
                                          area, ss);

  cudaDeviceSynchronize();
  float *pMin = (float *)d_vind_i;
  int *pIndx = (int *)(d_vind_i + 1);
  nppsMinIndx_32f(d_vind_f, ss, pMin, pIndx, (Npp8u *)d_vind_o);
  // min_ii<<<1, 1>>>(pIndx, d_vind, ss);

  // float fmin = vind[0];
  // for (unsigned i = 1; i < ss; i++) {
  //   if (vind[i] < fmin) {
  //     fmin = vind[i];
  //     idmin = i;
  //   }
  //   if (i == 10)
  //     printf("  PPP[%d] = %f  \n", 0, vind[0]);
  // }
  // // *idx = imin;
  // printf("  SSS  \n");

  unsigned idmin;
  float fmin;
  cudaMemcpy(&idmin, pIndx, sizeof(unsigned), cudaMemcpyDeviceToHost);
  cudaMemcpy(&fmin, pMin, sizeof(float), cudaMemcpyDeviceToHost);

  int ib = idmin;
  int iv = ib % (sv - 1);
  ib /= (sv - 1);
  int iw = ib % (sw - 1);
  ib /= (sw - 1);

  if (verbose)
    printf("========== Threshold [%d] [%d] [%d] - [%f]. ==========\n", iv, iw,
           ib, fmin);

  for (int t = 0; t < sz; t++)
    if (datav[t] > iv && dataw[t] > iw && datab[t] > ib)
      datao[t] = 1;
    else
      datao[t] = 0;

  cudaFree(d_datar);
  cudaFree(d_datab);
  cudaFree(d_dataw);
  cudaFree(d_datav);
  cudaFree(d_vind_i);
  cudaFree(d_vind_o);
  cudaFree(d_vind_f);
  // cudaFree(d_buff);
  delete vind;
}